# -*- coding: utf-8 -*-
import os
import geopandas as gpd
import rasterio
from rasterio.mask import mask
import numpy as np
from tqdm import tqdm
import earthaccess
from datetime import datetime

# ------------------------ 配置部分 ------------------------
START_DATE = "2024-01-01"
END_DATE = "2024-12-31"
BBOX = (117.4, 35.8, 118.4, 37.3)  # 淄博市经纬度范围
AREA_KM2 = 5965
TARGET_CITY = "淄博市"
SHAPEFILE_PATH = "data/shapefile/市.shp"
YEAR = datetime.strptime(START_DATE, "%Y-%m-%d").year
DOWNLOAD_DIR = f"data/modis/"
os.makedirs(DOWNLOAD_DIR, exist_ok=True)

# ------------------------ Earthdata 登录 ------------------------
if not os.path.exists(os.path.expanduser("~/.netrc")):
    raise FileNotFoundError("❌ 未找到 .netrc 文件，请参考 NASA Earthdata 配置说明生成。")

earthaccess.login(strategy="netrc")

# ------------------------ 数据搜索与下载 ------------------------
results = earthaccess.search_data(
    short_name="MOD16A2",
    cloud_hosted=True,
    bounding_box=BBOX,
    temporal=(START_DATE, END_DATE)
)

print(f"🔍 共找到 {len(results)} 个 granule，准备下载至：{DOWNLOAD_DIR}")
earthaccess.download(results, DOWNLOAD_DIR)

# ------------------------ 读取 shapefile ------------------------
gdf = gpd.read_file(SHAPEFILE_PATH)
if '市' not in gdf.columns:
    raise ValueError("❌ shapefile 中未找到 '市' 字段，请检查属性表字段名")

zibo = gdf[gdf['市'] == TARGET_CITY]
if zibo.empty:
    raise ValueError(f"❌ shapefile 中未找到名为 '{TARGET_CITY}' 的区域")

# ------------------------ 获取 MODIS 投影 ------------------------
hdf_files = [f for f in os.listdir(DOWNLOAD_DIR) if f.endswith(".hdf")]
if not hdf_files:
    raise FileNotFoundError("❌ 未找到任何下载的 HDF 文件")

first_hdf_path = os.path.join(DOWNLOAD_DIR, hdf_files[0]).replace("\\", "/")
et_layer = f'HDF4_EOS:EOS_GRID:"{first_hdf_path}":MOD_Grid_MOD16A2:ET_500m'

with rasterio.open(et_layer) as src:
    modis_crs = src.crs
    print("🧭 MODIS 数据投影:", modis_crs)

# ------------------------ 投影转换 ------------------------
zibo_sinu = zibo.to_crs(modis_crs)
print("📐 转换后的边界投影:", zibo_sinu.crs)

# ------------------------ 主函数：处理 ET 和 QC ------------------------


def read_and_mask_layer(hdf_path, subdataset_path, geometry):
    try:
        with rasterio.open(subdataset_path) as src:
            out_image, _ = mask(src, geometry, crop=True)
            return out_image[0].astype(np.float32)
    except Exception as e:
        print(
            f"⚠️ 读取失败：{os.path.basename(hdf_path)} → {subdataset_path}\n  错误详情: {e}")
        return None

# ------------------------ 子数据集名验证工具 ------------------------


def check_subdatasets(hdf_path):
    from osgeo import gdal
    hdf_path = os.path.abspath(hdf_path).replace("\\", "/")
    ds = gdal.Open(hdf_path)
    if ds is None:
        print(f"⚠️ 无法打开 HDF 文件: {hdf_path}")
        return
    print(f"\n📁 {os.path.basename(hdf_path)} 包含的数据集：")
    for name, desc in ds.GetSubDatasets():
        print(f" - {name}")


# ------------------------ 遍历 HDF 文件并处理 ------------------------
total_et = 0.0
count_success = 0
count_fail = 0
pixel_area_m2 = 500 * 500  # 500m分辨率对应的像素面积，单位平方米
total_et_volume = 0.0  # 用于累积蒸散发体积，单位立方米
for file in tqdm(hdf_files, desc="📦 正在处理 HDF 文件"):
    path = os.path.join(DOWNLOAD_DIR, file).replace("\\", "/")
    et_path = f'HDF4_EOS:EOS_GRID:"{path}":MOD_Grid_MOD16A2:ET_500m'
    qc_path = f'HDF4_EOS:EOS_GRID:"{path}":MOD_Grid_MOD16A2:ET_QC_500m'

    et_data = read_and_mask_layer(path, et_path, zibo_sinu.geometry)
    qc_data = read_and_mask_layer(path, qc_path, zibo_sinu.geometry)

    if et_data is None or qc_data is None:
        check_subdatasets(path)  # 如果出错就打印子数据集名帮助调试
        count_fail += 1
        continue

    # 替换无效值
    et_data[et_data == 32767] = np.nan
    qc_mask = (qc_data == 0) | (qc_data == 1)  # 保留高质量和次优像素
    et_data[~qc_mask] = np.nan

    # 计算体积：0.1mm单位换算为米，再乘以像素面积，得到体积m³
    et_m = et_data * 0.1 / 1000  # mm → m
    et_volume = et_m * pixel_area_m2

    valid_volume_sum = np.nansum(et_volume)

    if valid_volume_sum > 0:
        total_et_volume += valid_volume_sum
        count_success += 1
    else:
        print(f"⚠️ 有效蒸散发体积全部为无效：{file}")
        count_fail += 1

print(f"\n 成功处理：{count_success} 张影像， 失败：{count_fail} 张")
print(f" 2024年淄博市蒸散发总水量约为：{total_et_volume / 1e8:.2f} 亿立方米")  # 1亿立方米单位转换
# for file in tqdm(hdf_files, desc="📦 正在处理 HDF 文件"):
#     path = os.path.join(DOWNLOAD_DIR, file).replace("\\", "/")
#     et_path = f'HDF4_EOS:EOS_GRID:"{path}":MOD_Grid_MOD16A2:ET_500m'
#     qc_path = f'HDF4_EOS:EOS_GRID:"{path}":MOD_Grid_MOD16A2:ET_QC_500m'

#     et_data = read_and_mask_layer(path, et_path, zibo_sinu.geometry)
#     qc_data = read_and_mask_layer(path, qc_path, zibo_sinu.geometry)

#     if et_data is None or qc_data is None:
#         check_subdatasets(path)  # 如果出错就打印子数据集名帮助调试
#         count_fail += 1
#         continue

#     # 替换无效值
#     et_data[et_data == 32767] = np.nan
#     qc_mask = (qc_data == 0) | (qc_data == 1)  # 保留高质量和次优像素
#     et_data[~qc_mask] = np.nan

#     et_data *= 0.1  # 单位转换：0.1mm → mm
#     valid_sum = np.nansum(et_data)

#     if valid_sum > 0:
#         total_et += valid_sum
#         count_success += 1
#     else:
#         print(f"⚠️ 有效蒸散发值全部为无效：{file}")
#         count_fail += 1

# ------------------------ 输出结果 ------------------------
# print(f"\n✅ 成功处理：{count_success} 张影像，❌ 失败：{count_fail} 张")
# print(f"📏 蒸散发总量（mm）：{total_et:.2f}")
# volume_m3 = (total_et / 1000) * AREA_KM2 * 1e6
# print(f"💧 年蒸散发总水量约为：{volume_m3 / 1e8:.2f} 亿立方米")
